import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


# ==================== 数据处理部分（增强列名鲁棒性） ====================
def process_data():
    INPUT_FILE = "ori_message.xlsx"  # 原始数据路径
    OUTPUT_FILE = "monthly_average_data.xlsx"

    with pd.ExcelFile(INPUT_FILE) as excel_file:
        years = ['2016', '2017', '2018', '2019', '2020', '2021']
        all_monthly_data = []

        for year in years:
            # 读取工作表并标准化列名（去除首尾空格+统一半角符号）
            df = pd.read_excel(excel_file, sheet_name=year)
            df.columns = df.columns.str.strip()  # 去除列名首尾空格
            df.columns = df.columns.str.replace('（', '(').str.replace('）', ')')  # 全角转半角括号

            # 打印列名用于调试（确认是否匹配）
            print(f"工作表 {year} 的列名：{df.columns.tolist()}")

            # 检查关键列是否存在
            required_columns = ['年', '月', '流量(m3/s)', '含沙量(kg/m3)']
            missing_columns = [col for col in required_columns if col not in df.columns]
            if missing_columns:
                raise ValueError(f"工作表 {year} 缺少必要列：{missing_columns}")

            # 填充年月并转换类型
            df['年'] = df['年'].ffill().astype(int)
            df['月'] = df['月'].ffill().astype(int)

            # 按月分组计算（使用标准化后的列名）
            grouped = df.groupby(['年', '月']).agg({
                '流量(m3/s)': 'mean',
                '含沙量(kg/m3)': 'mean'
            }).reset_index()

            # 计算排沙量
            grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3)']
            grouped = grouped.round({
                '流量(m3/s)': 2,
                '含沙量(kg/m3)': 4,
                '排沙量(kg/s)': 2
            })

            all_monthly_data.append(grouped)

    # 合并并排序
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result.to_excel(OUTPUT_FILE, index=False)
    return result


# ==================== Mann-Kendall突变检验 ====================
def calculate_mk_statistics(series):
    x = np.asarray(series)
    n = len(x)

    # 计算正向统计量UFk
    uf = np.zeros(n)
    for k in range(n):
        counts = np.sum([np.sum(x[:i] < x[i]) for i in range(k + 1)])
        e_sk = k * (k + 1) / 4
        var_sk = k * (k + 1) * (2 * k + 5) / 72
        uf[k] = (counts - e_sk) / np.sqrt(var_sk) if var_sk != 0 else 0.0

    # 计算逆向统计量UBk（直接逆序计算）
    x_rev = x[::-1]
    uf_rev = np.zeros(n)
    for k in range(n):
        counts_rev = np.sum([np.sum(x_rev[:i] < x_rev[i]) for i in range(k + 1)])
        e_sk_rev = k * (k + 1) / 4
        var_sk_rev = k * (k + 1) * (2 * k + 5) / 72
        uf_rev[k] = (counts_rev - e_sk_rev) / np.sqrt(var_sk_rev) if var_sk_rev != 0 else 0.0

    ub = -uf_rev[::-1]
    ub[0] = 0.0
    return uf, ub


# ==================== 绘图函数 ====================
def plot_mk_test(processed_data):
    sediment_series = processed_data['排沙量(kg/s)']
    uf, ub = calculate_mk_statistics(sediment_series)
    months = np.arange(1, len(sediment_series) + 1)
    critical_value = 1.96

    plt.figure(figsize=(12, 6))
    plt.plot(months, uf, 'r-', label='UFk（正向趋势）', linewidth=1.5)
    plt.plot(months, ub, 'b-', label='UBk（逆向趋势）', linewidth=1.5)
    plt.axhline(critical_value, color='gray', linestyle='--', alpha=0.7, label='95%置信区间')
    plt.axhline(-critical_value, color='gray', linestyle='--', alpha=0.7)
    plt.title('2016-2021年每月平均排沙量M-K突变检验', fontsize=14, pad=20)
    plt.xlabel('时间（第1月=2016年1月，第72月=2021年12月）', fontsize=12)
    plt.ylabel('M-K统计量值', fontsize=12)
    plt.legend(loc='upper left', fontsize=10)
    plt.grid(alpha=0.1, linestyle='--')
    plt.tight_layout()
    plt.show()


# ==================== 主函数 ====================
if __name__ == "__main__":
    try:
        processed_data = process_data()
        print("\n数据处理成功，结果示例：")
        print(processed_data.head())  # 打印前5行验证数据

        plot_mk_test(processed_data)

    except ValueError as e:
        print(f"数据处理错误：{str(e)}")
    except Exception as e:
        print(f"程序运行失败，错误信息：{str(e)}")